#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2021 wanghuacheng <wanghuacheng@wanghuacheng-PC>
#
# Distributed under terms of the MIT license.

"""

"""
import sys
import os
import argparse
import pandas as pd
import talib as ta
import matplotlib.pyplot as plt
from talib.abstract import *
import datetime

date_range = [5, 13, 21, 34, 55, 89, 144, 233, 377]
len_date = len(date_range)

data_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data")
#merge_dir = os.path.join(data_dir, "history_k/merge/")
merge_dir = data_dir
merge_dates = os.listdir(merge_dir)
md = max(merge_dates)
merge_latest_dir = os.path.join(merge_dir, md)
code_names = os.listdir(merge_latest_dir)
#print(codes)
#sys.exit(0)


today = datetime.datetime.today()
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument('-c', '--code', type=str, default=None, help='code name')
parser.add_argument('-a', '--action', type=str, default="st1", help='code name')

ARGS = parser.parse_args()

today_str = today.strftime("%Y-%m-%d")


def gold_cross(last_diff, today_diff):
    return last_diff < 0 and today_diff >= 0

def dead_cross(last_diff, today_diff):
    return last_diff > 0 and today_diff <= 0

def get_data(code, startdate, enddate):
    rs = bs.query_history_k_data(code,
                                 "date,code,close,low,high,open,tradeStatus,peTTM,pbMRQ,psTTM,pcfNcfTTM",
                                 start_date=startdate, end_date=enddate,
                                 frequency="d", adjustflag="3")
    df = rs.get_data()
    # print(df.tail())
    return df

def is_top(df, index):
    close = df.at[index, 'close']
    dflag = True
    for i in range(1, 7):
        # print(df.loc[index, 'close'])
        if close >= df.at[index + i, 'close'] and close >= df.at[index - i, 'close']:
            continue
        else:
            dflag = False
            break
    return dflag

def is_bottom(df, index):
    close = df.at[index, 'close']
    dflag = True
    for i in range(1, 7):
        # print(df.loc[index, 'close'])
        if close <= df.at[index + i, 'close'] and close <= df.at[index - i, 'close']:
            continue
        else:
            dflag = False
            break
    return dflag


def find_st1(df):
    print(df.head())
    N = df.shape[0]
    if N < 500:
        return None

    features = []

    # close features
    sma_close_diff_names = []
    for d in date_range:
        diff_name = 'SMA_DIFF_' + str(d)
        sma_name = 'SMA_' + str(d)
        sma_close_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d)
        df[diff_name] = (df['close'] - df[sma_name]) / df[sma_name]


    sma_diff_names = []
    for i in range(0, len_date - 1):
        for j in range(i + 1, len_date):
            diff_name = "SMA_" + str(date_range[i]) + "_" + str(date_range[j]) + "_DIFF"
            sma_diff_names.append(diff_name)
            df[diff_name] = (df['SMA_' + str(date_range[i])] - df['SMA_' + str(date_range[j])]) / df['SMA_' + str(date_range[j])]

    sma_fes = ['SMA_' + str(d) for d in date_range]

    features = sma_close_diff_names + sma_diff_names
    print(features)


    # val features .............
    val_features = ["peTTM", "pbMRQ", "psTTM", "pcfNcfTTM"]
    val_diff_features = []
    for vf in val_features:
        df[vf] = df[vf].astype('float')
        sma_name = vf + '_377_SMA'
        df[sma_name] = SMA(df, timeperiod = 377, price = vf)
        # vf_mean = df[vf].mean()
        # print(vf_mean)
        vf_diff_name = sma_name + "_diff"
        val_diff_features.append(vf_diff_name)
        df[vf_diff_name] = (df[vf] - df[sma_name])/ df[sma_name]


    # volume
    sma_vol_diff_names = []
    for d in date_range:
        diff_name = 'VOL_SMA_DIFF_' + str(d)
        sma_name = 'VOL_SMA_' + str(d)
        sma_vol_diff_names.append(diff_name)

        df[sma_name] = SMA(df, timeperiod = d, price = 'volume')
        df[diff_name] = (df['volume'] - df[sma_name]) / df[sma_name]




    features = features + val_diff_features + sma_vol_diff_names

    for index, row in df.iterrows():
        if index < 500:
            continue
        if index > N - 10:
            break

        #print(row.date)
        dflag = is_bottom(df, index)
        tflag = is_top(df, index)
        if dflag:
            target = 0
            # print("BOTTOM:", row.date)
        if tflag:
            target = 1
            # print("TOP:", row.date)

        if dflag or tflag:
            # print(df.loc[index, features], target)
            fvals = (df.loc[index, features].values)
            yield fvals, target

def find_pe(df):
    #peMean = df['peTTM'].mean(skipna = True)
    pel = df.peTTM.astype('float')
    pe = (pel.mean())
    today_pe = pel[df.shape[0] - 1]
    print(today_pe, pe)
    # for index, row in df.iterrows():
        # print(index, row.peTTM)
    #print(peMean)


    #N = df.shape[0]
    #print(peMean, df.loc[N-1, 'peTTM'])

def find_macd(df):
    DIF = 1
    DEA = 2
    #df = df[df['tradeStatus'] == '1']
    if df.shape[0] < 50:
        return None
    macd_df = MACD(df, fastperiod=12, slowperiod=26, signalperiod=9)
    # print(macd_df)
    df = pd.concat([df, macd_df], axis = 1)
    N = int(df.shape[0])
    # print(df.shape)
    # dif 上穿 dea 白上穿黄
    # dif 下穿 dea 白下穿黄
    last_gold = None
    last_dead = None
    late_hist = df.at[N-1, 'macdhist']
    # print(df)
    for i in range(N - 1):
        day1_hist = df.at[i, 'macdhist']
        day2_hist = df.at[i + 1, 'macdhist']
        day2_dif =  df.at[i + 1, 'macd']
        dict1=df.loc[i + 1]
        dict2=df.loc[i]
        if gold_cross(day1_hist, day2_hist):
            # print(day2_dif > 0, "gold", dict2.to_json())
            last_gold = dict2['date']

        if dead_cross(day1_hist, day2_hist):
            # print(day2_dif > 0, "dead", dict2.to_json())
            last_dead = dict2['date']


    gddiff = get_date_diff(last_gold, last_dead)
    gdiff  = get_date_diff(today_str, last_gold)
    #print(last_gold, last_dead, late_hist)
    info = None
    if gdiff < 20 and gddiff > 0 and late_hist > 0:
        info = [last_dead, last_gold, gdiff]
    return info





def find_kdj(df):
    N = df.shape[0]
    matype = 0
    stoch = STOCH(df, fastk_period=9, slowk_matype=matype, slowk_period=3, slowd_period=3, slowd_matype=matype)
    df = pd.concat([df, stoch], axis=1)
    df['k_d'] = df['slowk'] - df['slowd']
    df['j'] = 3*df['slowk']- 2*df['slowd']

    last_dead = None
    last_gold = None
    last_kd = None
    for index, row in df.iterrows():
        if index < 50:
            continue
        last_kd = df.at[index - 1, 'k_d']
        today_kd = df.at[index, 'k_d']
        if gold_cross(last_kd, today_kd):
            last_gold = row.date
        elif dead_cross(last_kd, today_kd):
            last_dead = row.date

    gddiff = get_date_diff(last_gold, last_dead)
    gdiff  = get_date_diff(today_str, last_gold)
    info = None
    if gdiff < 20 and gddiff > 0 and last_kd > 0:
        info = [last_dead, last_gold, gdiff]

    return info



def find_rsi(df):
    N = df.shape[0]
    matype = 0
    rsis = []
    for d in [6, 12, 24]:
        dfn = 'rsi_' + str(d)
        df[dfn] = RSI(df, timeperiod = d)

        rsid = df.at[N-1 , dfn]
        rsis.append(rsid)

    score = 0
    # print(rsis)
    for index, r in enumerate(rsis):
        s = 0
        if r > 50 and r < 80:
            s = 1
        if r < 20:
            s = 1
        if r > 80:
            s = 1
        if r > 20 and r < 50:
            s = -1
        score += s * (index + 1)
    if score > 0:
        return str(score)
    else:
        return None

def find_bolling(df):
    N = df.shape[0]
    matype = 0
    rsis = []
    bands = BBANDS(df, timeperiod = 20, nbdevup = 2.0, nbdevdn = 2.0, matype = 0)
    df = pd.concat([df, bands], axis = 1)
    #upperband  middleband  lowerband
    # df.close - df.lowerband
    close = df.close.astype('float')
    df['BBP'] = (close - df.lowerband) / (df.upperband - df.lowerband)
    df['RSI'] = RSI(df, timeperiod = 14)
    x = df.loc[N-1, ['BBP', 'RSI']]
    print(x)


    score = 0
    print(df)
    if score > 0:
        return str(score)
    else:
        return None








def get_date_diff(d1, d2):
    dd1 = datetime.datetime.strptime(d1, '%Y-%m-%d')
    dd2 = datetime.datetime.strptime(d2, '%Y-%m-%d')
    return (dd1 - dd2).days

if __name__ == '__main__':
    code_df = pd.read_csv("./clist.csv")
    for index, row in code_df.iterrows():
        code = row['code']
        name = row['display_name']
        code_file = os.path.join(data_dir, code + "_" + name + ".csv")
        df = pd.read_csv(code_file)
        print(df.count)
